import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

data_tf = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize([0.5], [0.5])])

batch_size = 32
# 读取测试数据，train=True读取训练数据；train=False读取测试数据
train_dataset = datasets.MNIST(root='./data', train=True, transform=data_tf, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_tf)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

examples = enumerate(test_loader)  # img&label
batch_idx, (imgs, labels) = next(examples)  # 读取数据,batch_idx从0开始

print(labels)  # 读取标签数据
print(labels.shape)  # torch.Size([32])，因为batch_size为32

# -------------------------------数据显示--------------------------------------------
# 显示6张图片
import matplotlib.pyplot as plt

fig = plt.figure()
for i in range(6):
    plt.subplot(2, 3, i + 1)
    plt.tight_layout()
    plt.imshow(imgs[i][0], cmap='gray', interpolation='none')  # 子显示
    plt.title("Ground Truth: {}".format(labels[i]))  # 显示title
    plt.xticks([])
    plt.yticks([])

plt.show()